# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import numpy as np

import torch
import torch.nn as nn
from rsl_rl.utils import unpad_trajectories
from rsl_rl.modules.him_estimator import HIMEstimator

class ActorStudentTrain(nn.Module):
    is_recurrent = True

    def __init__(self,  num_actor_obs,
                 num_actions,
                 num_estimator,
                 actor_hidden_dims=[256, 256, 256],
                 hidden_size=0,
                 lstm_num_layers=0,
                 num_out_terrain = 0,
                 num_out_privilege = 0,
                 num_obs = 45,
                 activation='elu',
                 **kwargs):
        if kwargs:
            print("ActorCriticStudentOutput.__init__ got unexpected arguments, which will be ignored: " + str(
                [key for key in kwargs.keys()]))
        super(ActorStudentTrain, self).__init__()

        activation = get_activation(activation)

        if num_estimator is not None:
            self.estimator = HIMEstimator(temporal_steps=5, num_one_step_obs=45)
            print(f'Student him Estimator: {self.estimator.encoder}')

        mlp_input_dim_a = num_actor_obs

        if num_estimator is not None:
            self.estimator = HIMEstimator(temporal_steps=5, num_one_step_obs=45)
            print(f'Estimator: {self.estimator.encoder}')
        # Policy
        actor_layers = []
        actor_layers.append(nn.Linear(mlp_input_dim_a, actor_hidden_dims[0]))
        actor_layers.append(activation)
        for l in range(len(actor_hidden_dims)):
            if l == len(actor_hidden_dims) - 1:
                actor_layers.append(nn.Linear(actor_hidden_dims[l], num_actions))
            else:
                actor_layers.append(nn.Linear(actor_hidden_dims[l], actor_hidden_dims[l + 1]))
                actor_layers.append(activation)
        self.actor = nn.Sequential(*actor_layers)

        print(f"Student Actor MLP: {self.actor}")

        self.lstm = Memory(num_obs, lstm_num_layers, hidden_size)
        self.fc = nn.Linear(hidden_size, num_out_terrain + num_out_privilege)
        self.num_out_terrain = num_out_terrain
        self.num_out_privilege = num_out_privilege
        print(f"LSTM: {self.lstm}")
        print(f"LSTM_fc: {self.fc}")

    def forward(self, observations):
        lstm_out = self.lstm(observations).squeeze(0)
        output = self.fc(lstm_out)
        him_latent = self.estimator(observations)
        terrain_latent = output[:, :self.num_out_terrain]
        privilege_latent = output[:, -self.num_out_privilege:]
        actions_mean = self.actor(torch.cat((observations, him_latent, terrain_latent, privilege_latent), dim=-1))
        return actions_mean


    def student_forward(self, observations):
        lstm_out = self.lstm(observations).squeeze(0)
        output = self.fc(lstm_out)
        terrain_latent = output[:, :self.num_out_terrain]
        privilege_latent = output[:, -self.num_out_privilege:]
        actions_mean = self.actor(torch.cat((observations, terrain_latent, privilege_latent), dim=-1))
        return actions_mean

    def student_act(self, observations, masks=None, hidden_states=None):
        if hidden_states is not None:
            lstm_out = self.lstm(observations, masks, hidden_states).squeeze(0)
            output = self.fc(lstm_out).reshape(-1, self.num_out_privilege + self.num_out_terrain)
        else:
            lstm_out = self.lstm(observations).squeeze(0)
            output = self.fc(lstm_out)
        if masks is not None:
            observations = unpad_trajectories(observations,masks).reshape(-1,observations.shape[-1])
        if hasattr(self, 'estimator'):
            with torch.no_grad():
                him_latent = self.estimator(observations)
        terrain_latent = output[:, :self.num_out_terrain]
        privilege_latent = output[:, -self.num_out_privilege:]
        mean = self.actor(torch.cat((observations, him_latent, terrain_latent, privilege_latent), dim=-1))
        return output, mean

    def student_act_inference(self, observations, masks=None, hidden_states=None):
        if hidden_states is not None:
            lstm_out = self.lstm(observations, masks, hidden_states).squeeze(0)
            output = self.fc(lstm_out).reshape(-1, self.num_out_privilege + self.num_out_terrain)
        else:
            lstm_out = self.lstm(observations).squeeze(0)
            output = self.fc(lstm_out)
        if masks is not None:
            observations = unpad_trajectories(observations,masks).reshape(-1,observations.shape[-1])
        if hasattr(self, 'estimator'):
            with torch.no_grad():
                him_latent = self.estimator(observations)
                terrain_latent = output[:, :self.num_out_terrain]
                privilege_latent = output[:, -self.num_out_privilege:]
        mean = self.actor(torch.cat((observations, him_latent, terrain_latent, privilege_latent), dim=-1))
        return mean.detach()

class Memory(torch.nn.Module):
    def __init__(self, input_size, num_layers=1, hidden_size=256):
        super().__init__()
        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers)
        self.hidden_states = None
        self.num_layers = num_layers
        self.hidden_size = hidden_size

    def forward(self, input, masks=None, hidden_states=None):
        batch_mode = masks is not None
        if batch_mode:
            # batch mode (policy update): need saved hidden states
            if hidden_states is None:
                raise ValueError("Hidden states not passed to memory module during policy update")
            out, _ = self.lstm(input, hidden_states)
            out = unpad_trajectories(out, masks)
        else:
            # inference mode (collection): use hidden states of last step
            out, self.hidden_states = self.lstm(input.unsqueeze(0), self.hidden_states)
        return out

    def reset(self, dones=None):
        # When the RNN is an LSTM, self.hidden_states_a is a list with hidden_state and cell_state
        for hidden_state in self.hidden_states:
            hidden_state[..., dones, :] = 0.0

def get_activation(act_name):
    if act_name == "elu":
        return nn.ELU()
    elif act_name == "selu":
        return nn.SELU()
    elif act_name == "relu":
        return nn.ReLU()
    elif act_name == "crelu":
        return nn.ReLU()
    elif act_name == "lrelu":
        return nn.LeakyReLU()
    elif act_name == "tanh":
        return nn.Tanh()
    elif act_name == "sigmoid":
        return nn.Sigmoid()
    else:
        print("invalid activation function!")
        return None
